package com.yunyi.handler;

import com.yunyi.constants.RoleConstants;
import com.yunyi.entity.mongo.ChatMessageDocument;
import com.yunyi.mongo.ChatHistoryRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;

@RequiredArgsConstructor
@Component
public class MongoChatMemoryHandler implements ChatMemory {


    private final ChatHistoryRepository chatHistoryRepository;


    @Override
    public void add(String conversationId, List<Message> messages) {
        List<ChatMessageDocument> list = new ArrayList<>();
        for (Message message : messages) {
            ChatMessageDocument document = new ChatMessageDocument();
            if (message instanceof UserMessage) {
                document.setRole(RoleConstants.USER_ROLE_CONSTANTS);
            } else if (message instanceof SystemMessage) {
                document.setRole(RoleConstants.SYSTEM_ROLE_CONSTANTS);
            } else if (message instanceof ToolResponseMessage) {
                document.setRole(RoleConstants.TOOL_ROLE_CONSTANTS);
            } else if (message instanceof AssistantMessage) {
                document.setRole(RoleConstants.ASSISTANT_ROLE_CONSTANTS);
            }
            document.setSessionId(conversationId);
            document.setContent(message.getText());
            document.setCreateTime(LocalDateTime.now());
            list.add(document);
        }
        chatHistoryRepository.saveAll(list);
    }

    @Override
    public List<Message> get(String conversationId) {
        List<ChatMessageDocument> chatMessages = chatHistoryRepository.getChatMessageDocumentBySessionId(conversationId)
                .stream()
                .sorted((o1, o2) -> o2.getCreateTime().compareTo(o1.getCreateTime()))
                .toList();
        List<Message> chatList = new ArrayList<>();
        for (ChatMessageDocument chatMessage : chatMessages) {
            switch (chatMessage.getRole()) {
                case RoleConstants.USER_ROLE_CONSTANTS:
                    chatList.add(new UserMessage(chatMessage.getContent()));
                    break;
                case RoleConstants.ASSISTANT_ROLE_CONSTANTS:
                    chatList.add(new AssistantMessage(chatMessage.getContent()));
                    break;
                case RoleConstants.SYSTEM_ROLE_CONSTANTS:
                    chatList.add(new SystemMessage(chatMessage.getContent()));
                    break;
                default:
                    throw new IllegalArgumentException("未知的消息类型");
            }
        }

        return chatList;
    }

    @Override
    public void clear(String conversationId) {
        chatHistoryRepository.deleteChatMessageDocumentBySessionId(conversationId);
    }
}
